# -*- coding: utf-8 -*-
# !/usr/bin/python3
"""
Author :      wu
Description :
"""

import tensorflow as tf
from tensorflow.keras import layers, models, regularizers, constraints
import pandas as pd

from feature_column import prepare_df_data, df_to_dataset, build_feature_columns


# 直接定义损失函数
def focal_loss(gamma=2.0, alpha=0.75):

    def focal_loss_fixed(y_true, y_pre):
        bce = tf.losses.binary_crossentropy(y_true, y_pre)
        p_t = (float(y_true) * y_pre) + (float(1 - y_true) * float(1 - y_pre))
        alpha_factor = float(y_true) * alpha + float(1 - y_true) * float(1 - alpha)
        modulating_factor = tf.pow(1.0 - p_t, gamma)
        loss = tf.reduce_sum(bce * alpha_factor * modulating_factor, axis=-1)
        return loss

    return focal_loss_fixed


# 继承tf的Loss类
class FocalLoss(tf.keras.losses.Loss):
    def __init__(self, gamma=2.0, alpha=0.75, name="focal_loss"):
        super(FocalLoss, self).__init__(name=name)
        self.gamma = gamma
        self.alpha = alpha

    def call(self, y_true, y_pre):
        bce = tf.losses.binary_crossentropy(y_true, y_pre)
        p_t = (float(y_true) * float(y_pre)) + (float(1 - y_true) * float(1 - y_pre))
        alpha_factor = float(y_true) * self.alpha + float((1 - y_true)) * (1 - self.alpha)
        modulating_factor = tf.pow(1.0 - p_t, self.gamma)
        loss = tf.reduce_sum(bce * alpha_factor * modulating_factor, axis=-1)

        return loss


def main():
    tf.keras.backend.clear_session()

    df = pd.read_csv("../data/titanic/train.csv")
    df = df.sample(frac=1).reset_index(drop=True)  # 打乱顺序

    df_train_raw = df.iloc[:int(len(df) * 0.7), :]
    df_test_raw = df.iloc[int(len(df) * 0.7):, :]

    df_raw = pd.concat([df_train_raw, df_test_raw])
    df_raw = prepare_df_data(df_raw)
    df_train = df_raw.iloc[:len(df_train_raw), :]
    df_test = df_raw.iloc[len(df_train_raw):, :]

    ds_train = df_to_dataset(df_train)
    ds_test = df_to_dataset(df_test)

    feature_columns = build_feature_columns(df_raw)

    tf.keras.backend.clear_session()
    model = tf.keras.Sequential([
        layers.DenseFeatures(feature_columns),
        layers.Dense(64, activation="relu"),
        layers.Dense(32, activation="relu"),
        layers.Dense(1, activation="sigmoid")
    ])

    model.compile(optimizer="Adam", loss=focal_loss(), metrics=["accuracy"])
    history = model.fit(ds_train, validation_data=ds_test, epochs=5)
    model.summary()
    print("loss: {}".format(model.loss))


if __name__ == "__main__":
    main()
